import net from 'net';
import fs from 'fs';
import path from 'path';
import os from 'os';
import { fork } from 'child_process';
import { fileURLToPath } from 'url';

const PORT = 5000;
const SAMPLE_RATE = 16000;
const CHANNELS = 1;
const BIT_DEPTH = 16;
const OUTPUT_DIR = './temp_audio';
const SEGMENT_BYTES = 1024 * 512;

if (!fs.existsSync(OUTPUT_DIR)) {
    fs.mkdirSync(OUTPUT_DIR, { recursive: true });
}

// 子进程写 WAV
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
const workerPath = path.join(__dirname, 'wav_worker.js');
const worker = fork(workerPath, [], { stdio: ['ignore', 'inherit', 'inherit', 'ipc'] });
let workerExited = false;
let isShuttingDown = false;

worker.on('message', (msg) => {
    if (msg?.type === 'exited') workerExited = true;
});
worker.on('exit', () => { workerExited = true; });

worker.send({
    type: 'init',
    config: { sampleRate: SAMPLE_RATE, channels: CHANNELS, bitDepth: BIT_DEPTH, outputDir: OUTPUT_DIR }
});

const sockets = new Set();
let fileSeq = 0;

const server = net.createServer((socket) => {
    console.log('TCP 客户端已连接:', `${socket.remoteAddress}:${socket.remotePort}`);
    sockets.add(socket);

    // 每个连接的上下文
    const ctx = {
        identified: false,     // 是否已经握手完成
        mac: null,
        chipid: null,
        pendingBuffers: [],
        pendingBytes: 0,
        total: 0
    };
    socket._ctx = ctx;

    // 发送到子进程
    function sendToWorker(pcm) {
        const safeMac = ctx.mac ? ctx.mac.replace(/:/g, '-') : 'unknown';
        const name = `${safeMac}_${Date.now()}_${++fileSeq}.wav`;

        if (isShuttingDown || workerExited || !worker?.connected) {
            console.warn('子进程不可用，直接主进程保存');
            const fallback = path.join(OUTPUT_DIR, `fallback_${name}`);
            saveWavSync(fallback, pcm);
            return;
        }
        worker.send({ type: 'segment', pcm, name });
    }

    function saveWavSync(filepath, pcm) {
        try {
            const dataLength = pcm.length;
            const fileLength = 36 + dataLength;
            const header = Buffer.alloc(44);
            header.write('RIFF', 0);
            header.writeUInt32LE(fileLength, 4);
            header.write('WAVE', 8);
            header.write('fmt ', 12);
            header.writeUInt32LE(16, 16);
            header.writeUInt16LE(1, 20);
            header.writeUInt16LE(CHANNELS, 22);
            header.writeUInt32LE(SAMPLE_RATE, 24);
            header.writeUInt32LE(SAMPLE_RATE * CHANNELS * (BIT_DEPTH / 8), 28);
            header.writeUInt16LE(CHANNELS * (BIT_DEPTH / 8), 32);
            header.writeUInt16LE(BIT_DEPTH, 34);
            header.write('data', 36);
            header.writeUInt32LE(dataLength, 40);
            const wav = Buffer.concat([header, pcm]);
            fs.writeFileSync(filepath, wav);
        } catch (e) {
            console.error('保存 WAV 失败:', e);
        }
    }

    function writeSegment(bytesToWrite) {
        let remaining = bytesToWrite;
        const parts = [];
        while (remaining > 0 && ctx.pendingBuffers.length > 0) {
            const first = ctx.pendingBuffers[0];
            if (first.length <= remaining) {
                parts.push(first);
                ctx.pendingBuffers.shift();
                ctx.pendingBytes -= first.length;
                remaining -= first.length;
            } else {
                const slice = first.subarray(0, remaining);
                const rest = first.subarray(remaining);
                parts.push(slice);
                ctx.pendingBuffers[0] = rest;
                ctx.pendingBytes -= remaining;
                remaining = 0;
            }
        }
        if (parts.length > 0) {
            const pcm = Buffer.concat(parts);
            sendToWorker(pcm);
        }
    }

    socket.on('data', (buf) => {
        if (!ctx.identified) {
            // 第一次收到数据当握手信息
            const msg = buf.toString().trim();
            try {
                const obj = JSON.parse(msg);
                ctx.mac = obj.mac;
                ctx.chipid = obj.chipid;
                ctx.identified = true;
                console.log(`客户端标识成功: MAC=${ctx.mac}, ChipID=${ctx.chipid}`);
                socket.write("ok\n"); // 发送握手确认
            } catch (e) {
                console.error('握手信息解析失败:', e?.message || e);
                socket.destroy();
            }
            return;
        }

        // 已经握手，开始接收音频数据
        ctx.pendingBuffers.push(buf);
        ctx.pendingBytes += buf.length;
        ctx.total += buf.length;

        if (ctx.pendingBytes >= SEGMENT_BYTES) {
            while (ctx.pendingBytes >= SEGMENT_BYTES) {
                writeSegment(SEGMENT_BYTES);
            }
        }
    });

    socket.on('end', () => {
        console.log(`客户端 ${ctx.mac || 'unknown'} 断开，总字节:`, ctx.total);
        if (ctx.pendingBytes > 0) writeSegment(ctx.pendingBytes);
    });

    socket.on('close', () => sockets.delete(socket));
    socket.on('error', (err) => console.error('TCP 错误:', err?.message || err));
});

server.listen(PORT, () => {
    console.log(`TCP 服务器已启动，监听端口 ${PORT}`);
    const nets = os.networkInterfaces();
    for (const name of Object.keys(nets)) {
        for (const n of nets[name] || []) {
            if (n.family === 'IPv4' && !n.internal) console.log(`  ${n.address}:${PORT}`);
        }
    }
});

process.on('SIGINT', () => {
    console.log('\n正在停止 TCP 服务器...');
    isShuttingDown = true;
    try { server.close(); } catch { }
    for (const s of sockets) {
        try { if (s._ctx.pendingBytes > 0) writeSegment(s._ctx.pendingBytes); } catch { }
        try { s.destroy(); } catch { }
    }
    if (!workerExited && worker?.connected) {
        try { worker.send({ type: 'exit' }); } catch { }
    }
    setTimeout(() => process.exit(0), 5000).unref();
});
